Skip to content

perf(index): parallelize kmeans centroid recomputation - #8560

Open
isaac-dasari wants to merge 3 commits into
lance-format:mainfrom
isaac-dasari:perf/kmeans-centroid-recompute-6369
Open

perf(index): parallelize kmeans centroid recomputation#8560
isaac-dasari wants to merge 3 commits into
lance-format:mainfrom
isaac-dasari:perf/kmeans-centroid-recompute-6369

Conversation

@isaac-dasari

Copy link
Copy Markdown

Summary

  • recompute KMeans centroids by partitioning input vectors across workers so each input vector is scanned once
  • merge thread-local centroid accumulators in parallel and cap their aggregate allocation at 64 MiB
  • add correctness coverage for f16, f32, f64, ignored memberships, empty-cluster splitting, and the memory cap
  • add a focused Criterion benchmark and include it in the existing Rust benchmark workflow

The previous implementation divided centroids among workers, but every worker scanned the complete input. That made centroid recomputation perform O(num_vectors * workers) input reads. This changes the input scan to O(num_vectors) while avoiding shared-write contention.

Closes #6369.

Validation

Passed locally:

  • cargo fmt --all -- --check
  • cargo check -p lance-index --tests --benches --profile bench -j 1
  • cargo clippy -p lance-index --tests --benches --profile bench -j 1 -- -D warnings
  • git diff --check

The optimized test/benchmark linker exceeded the available resources on the local 7.6 GiB development machine, so executable test and benchmark runs are intentionally delegated to GitHub Actions. The PR is draft pending those results.

@github-actions github-actions Bot added A-index Vector index, linalg, tokenizer A-ci CI / build workflows performance labels Aug 16, 2026
@isaac-dasari
isaac-dasari marked this pull request as ready for review August 17, 2026 07:00
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 17, 2026
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 17, 2026
lance-gatekeeper[bot]

This comment was marked as outdated.

@lance-gatekeeper lance-gatekeeper Bot added the K-changes Latest Gatekeeper recommendation requests changes. label Aug 17, 2026
@lance-gatekeeper lance-gatekeeper Bot removed the K-changes Latest Gatekeeper recommendation requests changes. label Aug 20, 2026

@lance-gatekeeper lance-gatekeeper Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Gate recommendation: approve.

The revision replaces dense per-worker accumulators with a stable membership-to-row index, preserving per-cluster accumulation order while removing repeated membership scans. The previously reproduced low-sample, default high-dimensional, and large-grid cases now match or beat base without reintroducing the large accumulator footprint.

@lance-gatekeeper lance-gatekeeper Bot added the K-approved Latest Gatekeeper recommendation permits acceptance. label Aug 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

A-ci CI / build workflows A-index Vector index, linalg, tokenizer K-approved Latest Gatekeeper recommendation permits acceptance. performance

Projects

None yet

Development

Successfully merging this pull request may close these issues.

KMeans to_kmeans centroid recomputation is suboptimal — each core redundantly scans all data

1 participant